#include "sgEigenSpace.h"
#include <iostream>
#include <mkl_spblas.h>

using namespace SG::Algebra;

//=============================================================================
// EigenVector

template <typename TDataType>
EigenVector<TDataType>::EigenVector ()
{}

template <typename TDataType>
EigenVector<TDataType>::EigenVector (Sparsity *sp)
{}

template <typename TDataType>
EigenVector<TDataType>::EigenVector (const Index_t &n)
{}

template <typename TDataType>
EigenVector<TDataType>::EigenVector (Eigen::VectorX<TDataType> &x)
{}

template <typename TDataType>
EigenVector<TDataType>::EigenVector (const EigenVector<TDataType> &x)
{}

template <typename TDataType>
EigenVector<TDataType>::EigenVector (EigenVector<TDataType> &&x)
{}

template <typename TDataType>
EigenVector<TDataType>::~EigenVector ()
{}

template <typename TDataType>
EigenVector<TDataType> &EigenVector<TDataType>::operator= (const EigenVector<TDataType> &x)
{
    return *this;
}

template <typename TDataType>
Space<TDataType> *EigenVector<TDataType>::getSpace ()
{
    static EigenSpace<TDataType> eigen_space;
    return &eigen_space;
}

template <typename TDataType>
void EigenVector<TDataType>::close ()
{}

template <typename TDataType>
void EigenVector<TDataType>::clear ()
{}

template <typename TDataType>
void EigenVector<TDataType>::beginAssembly ()
{}

template <typename TDataType>
void EigenVector<TDataType>::endAssembly ()
{}

template <typename TDataType>
void EigenVector<TDataType>::assemble ()
{}

// template <typename TDataType>
// bool EigenVector<TDataType>::assembled ()
// {
//     return true;
// }

template <typename TDataType>
void EigenVector<TDataType>::resize (Index_t size)
{
    Eigen::VectorX<TDataType>::resize (size);
}

template <typename TDataType>
void EigenVector<TDataType>::conservativeResize (Index_t size)
{
    Eigen::VectorX<TDataType>::conservativeResize (size);
}

template <typename TDataType>
Index_t EigenVector<TDataType>::rows () const
{
    return Eigen::VectorX<TDataType>::rows ();
}

template <typename TDataType>
Index_t EigenVector<TDataType>::row_start () const
{
    return 0;
}

template <typename TDataType>
Index_t EigenVector<TDataType>::row_end () const
{
    return Eigen::VectorX<TDataType>::rows ();
}

template <typename TDataType>
void EigenVector<TDataType>::setZero ()
{
    Eigen::VectorX<TDataType>::setZero ();
}

template <typename TDataType>
void EigenVector<TDataType>::set (const Index_t &i, const TDataType &value)
{
    this->coeffRef (i) = value;
}

template <typename TDataType>
void EigenVector<TDataType>::add (const Index_t &i, const TDataType &value)
{
    this->coeffRef (i) += value;
}

template <typename TDataType>
TDataType EigenVector<TDataType>::operator() (const Index_t &i) const
{
    return this->coeff (i);
}

template <typename TDataType>
TDataType EigenVector<TDataType>::operator[] (const Index_t &i) const
{
    return this->coeff (i);
}

// template <typename TDataType>
// TDataType EigenVector<TDataType>::dot(const EigenVector<TDataType> &other) const
// {
//     return 0.0;
// }

template <typename TDataType>
void EigenVector<TDataType>::print (std::ostream &out) const
{
    out << *this;
}

template <typename TDataType>
void EigenVector<TDataType>::create_subvector (Vector<TDataType> &subvector, const GeneralIndexSet &rows) const
{

}

template <typename TDataType>
void EigenVector<TDataType>::create_subvector (Vector<TDataType> &subvector, const StrideIndexSet &rows) const
{
    
}

template <typename TDataType>
void EigenVector<TDataType>::max (Index_t &p, TDataType &val) const
{

}

template <typename TDataType>
void EigenVector<TDataType>::linfty_norm (TDataType &value) const
{

}

//=============================================================================
// EigenDenseMatrix

template <typename TDataType>
void EigenDenseMatrix<TDataType>::print (std::ostream &out) const
{
    // out << *this;
}

//=============================================================================
// EigenSparseMatrix

template <typename TDataType>
EigenSparseMatrix<TDataType>::EigenSparseMatrix ()
{}

template <typename TDataType>
EigenSparseMatrix<TDataType>::EigenSparseMatrix (Sparsity *sp)
{}

template <typename TDataType>
EigenSparseMatrix<TDataType>::EigenSparseMatrix (Index_t M, Index_t N)
{}

template <typename TDataType>
EigenSparseMatrix<TDataType>::EigenSparseMatrix (Eigen::SparseMatrix<TDataType> &A)
{}

template <typename TDataType>
EigenSparseMatrix<TDataType>::EigenSparseMatrix (const EigenSparseMatrix<TDataType> &A)
{}

template <typename TDataType>
EigenSparseMatrix<TDataType>::EigenSparseMatrix (EigenSparseMatrix<TDataType> &&A)
{}

template <typename TDataType>
EigenSparseMatrix<TDataType>::~EigenSparseMatrix ()
{}

template <typename TDataType>
EigenSparseMatrix<TDataType> &EigenSparseMatrix<TDataType>::operator= (const EigenSparseMatrix<TDataType> &A)
{
    return *this;
}

template <typename TDataType>
Space<TDataType> *EigenSparseMatrix<TDataType>::getSpace ()
{
    static EigenSpace<TDataType> eigen_space;
    return &eigen_space;
}

template <typename TDataType>
void EigenSparseMatrix<TDataType>::clear ()
{

}

template <typename TDataType>
void EigenSparseMatrix<TDataType>::close ()
{}

template <typename TDataType>
void EigenSparseMatrix<TDataType>::beginAssembly ()
{}

template <typename TDataType>
void EigenSparseMatrix<TDataType>::endAssembly ()
{
    mat.makeCompressed ();
}

template <typename TDataType>
void EigenSparseMatrix<TDataType>::assemble ()
{
    mat.makeCompressed ();
}

template <typename TDataType>
bool EigenSparseMatrix<TDataType>::assembled () const
{
    return mat.isCompressed ();
}

template <typename TDataType>
void EigenSparseMatrix<TDataType>::resize (Index_t rows, Index_t cols)
{}

template <typename TDataType>
void EigenSparseMatrix<TDataType>::conservativeResize (Index_t rows, Index_t cols)
{}

template <typename TDataType>
Index_t EigenSparseMatrix<TDataType>::rows () const
{
    return mat.rows ();
}

template <typename TDataType>
Index_t EigenSparseMatrix<TDataType>::cols () const
{
    return mat.cols ();
}

template <typename TDataType>
Index_t EigenSparseMatrix<TDataType>::row_start () const
{
    return 0;
}

template <typename TDataType>
Index_t EigenSparseMatrix<TDataType>::row_end () const
{
    return mat.rows ();
}

template <typename TDataType>
Index_t EigenSparseMatrix<TDataType>::col_start () const
{
    return 0;
}

template <typename TDataType>
Index_t EigenSparseMatrix<TDataType>::col_end () const
{
    return mat.cols ();
}

template <typename TDataType>
void EigenSparseMatrix<TDataType>::setZero ()
{
    mat.setZero ();
}

template <typename TDataType>
void EigenSparseMatrix<TDataType>::setIdentity ()
{
    mat.setIdentity ();
}

template <typename TDataType>
void EigenSparseMatrix<TDataType>::set (const Index_t &i, const Index_t &j, const TDataType &value)
{
    mat.coeffRef (i, j) = value;
}

template <typename TDataType>
void EigenSparseMatrix<TDataType>::add (const Index_t &i, const Index_t &j, const TDataType &value)
{
    mat.coeffRef (i, j) += value;
}

// template <typename TDataType>
// void EigenSparseMatrix<TDataType>::add(const DenseMatrix<TDataType> &dm, const IndexMap_t &rows, const IndexMap_t &cols)
// {
// }

// template <typename TDataType>
// void EigenSparseMatrix<TDataType>::add(const SparseMatrix<TDataType> & /*spm*/, const IndexMap_t & /*rows*/, const IndexMap_t & /*cols*/)
// {

// }

template <typename TDataType>
TDataType EigenSparseMatrix<TDataType>::operator() (const Index_t &i, const Index_t &j) const
{
    return mat.coeff (i, j);
}

template <typename TDataType>
Index_t EigenSparseMatrix<TDataType>::nonZeros () const
{
    return mat.nonZeros ();
}

template <typename TDataType>
void EigenSparseMatrix<TDataType>::get_csr (std::vector<Index_t> &ia, std::vector<Index_t> &ja, std::vector<TDataType> &a) const
{}

template <typename TDataType>
void EigenSparseMatrix<TDataType>::print (std::ostream &out) const
{
    // out << mat;
}

template <typename TDataType>
void EigenSparseMatrix<TDataType>::create_submatrix (SparseMatrix<TDataType> &submatrix,
                                                     const GeneralIndexSet   &rows,
                                                     const GeneralIndexSet   &cols) const
{}

template <typename TDataType>
void EigenSparseMatrix<TDataType>::create_submatrix (SparseMatrix<TDataType> &submatrix,
                                                     const StrideIndexSet    &rows,
                                                     const StrideIndexSet    &cols) const
{}

template <typename TDataType>
void EigenSparseMatrix<TDataType>::create_submatrix_nosort (SparseMatrix<TDataType> & /*submatrix*/,
                                                            const GeneralIndexSet & /*rows*/,
                                                            const GeneralIndexSet & /*cols*/) const
{}

template <typename TDataType>
void EigenSparseMatrix<TDataType>::create_submatrix_nosort (SparseMatrix<TDataType> &submatrix,
                                                            const StrideIndexSet    &rows,
                                                            const StrideIndexSet    &cols) const
{}

template <typename TDataType>
void EigenSparseMatrix<TDataType>::get_diagonal (Vector<TDataType> &dest) const
{

}

template <typename TDataType>
void EigenSparseMatrix<TDataType>::get_transpose (SparseMatrix<TDataType> &dest) const
{}

template <typename TDataType>
void EigenSparseMatrix<TDataType>::get_inverse (SparseMatrix<TDataType> &dest) const
{

}

template <typename TDataType>
void EigenSparseMatrix<TDataType>::scale (const TDataType scale)
{

}

template <typename TDataType>
void EigenSparseMatrix<TDataType>::matrix_matrix_mult (SparseMatrix<TDataType> &X, SparseMatrix<TDataType> &Y)
{

}

//=============================================================================
// Operators

SG_NAMESPACE_OPEN 
namespace Algebra
{
    template <typename TDataType>
    EigenVector<TDataType> operator+ (const EigenVector<TDataType> &x, const EigenVector<TDataType> &y)
    {
        return EigenVector<TDataType> ();
    }

    template <typename TDataType>
    EigenVector<TDataType> operator- (const EigenVector<TDataType> &x, const EigenVector<TDataType> &y)
    {
        return EigenVector<TDataType> ();
    }

    template <typename TDataType>
    EigenVector<TDataType> operator* (const TDataType &s, const EigenVector<TDataType> &x)
    {
        return EigenVector<TDataType> ();
    }

    template <typename TDataType>
    TDataType operator* (const EigenVector<TDataType> &x, const EigenVector<TDataType> &y)
    {
        return 0.0;
    }

    template <typename TDataType>
    EigenVector<TDataType> operator/ (const EigenVector<TDataType> &x, const TDataType &s)
    {
        return EigenVector<TDataType> ();
    }

    template <typename TDataType>
    EigenSparseMatrix<TDataType> operator+ (const EigenSparseMatrix<TDataType> &X, const EigenSparseMatrix<TDataType> &Y)
    {
        return EigenSparseMatrix<TDataType> ();
    }

    template <typename TDataType>
    EigenSparseMatrix<TDataType> operator- (const EigenSparseMatrix<TDataType> &X, const EigenSparseMatrix<TDataType> &Y)
    {
        return EigenSparseMatrix<TDataType> ();
    }

    template <typename TDataType>
    EigenSparseMatrix<TDataType> operator* (const TDataType &s, const EigenSparseMatrix<TDataType> &A)
    {
        return EigenSparseMatrix<TDataType> ();
    }

    template <typename TDataType>
    EigenVector<TDataType> operator* (const EigenSparseMatrix<TDataType> &A, const EigenVector<TDataType> &x)
    {
        return EigenVector<TDataType> ();
    }

    template <typename TDataType>
    TDataType quadratic (const EigenSparseMatrix<TDataType> &A, const EigenVector<TDataType> &x)
    {
        return 0.0;
    }

    template <typename TDataType>
    std::ostream &operator<< (std::ostream &out, const EigenVector<TDataType> &v)
    {
        v.print (out);
        return out;
    }

    template <typename TDataType>
    std::ostream &operator<< (std::ostream &out, const EigenDenseMatrix<TDataType> &m)
    {
        m.print (out);
        return out;
    }

    template <typename TDataType>
    std::ostream &operator<< (std::ostream &out, const EigenSparseMatrix<TDataType> &m)
    {
        m.print (out);
        return out;
    }
}  // namespace Algebra
SG_NAMESPACE_CLOSE

//=============================================================================
// Explicit instantiations

SG_NAMESPACE_OPEN
namespace Algebra
{
    //------------------------------------------------------------------
    // Real_t
    template class ALGEBRA_EXPORT EigenVector<Real_t>;
    template class ALGEBRA_EXPORT EigenDenseMatrix<Real_t>;
    template class ALGEBRA_EXPORT EigenSparseMatrix<Real_t>;

    template ALGEBRA_EXPORT        EigenVector<Real_t> operator+ <Real_t> (const EigenVector<Real_t> &x, const EigenVector<Real_t> &y);
    template ALGEBRA_EXPORT        EigenVector<Real_t> operator- <Real_t> (const EigenVector<Real_t> &, const EigenVector<Real_t> &);
    template ALGEBRA_EXPORT        EigenVector<Real_t> operator* <Real_t> (const Real_t &s, const EigenVector<Real_t> &x);
    template ALGEBRA_EXPORT Real_t operator* <Real_t> (const EigenVector<Real_t> &x, const EigenVector<Real_t> &y);
    template ALGEBRA_EXPORT        EigenVector<Real_t> operator/ <Real_t> (const EigenVector<Real_t> &x, const Real_t &s);

    template ALGEBRA_EXPORT EigenSparseMatrix<Real_t> operator+ <Real_t> (const EigenSparseMatrix<Real_t> &,
                                                                             const EigenSparseMatrix<Real_t> &);
    template ALGEBRA_EXPORT EigenSparseMatrix<Real_t> operator- <Real_t> (const EigenSparseMatrix<Real_t> &,
                                                                             const EigenSparseMatrix<Real_t> &);
    template ALGEBRA_EXPORT EigenSparseMatrix<Real_t> operator* <Real_t> (const Real_t &, const EigenSparseMatrix<Real_t> &);
    template ALGEBRA_EXPORT EigenVector<Real_t> operator* <Real_t> (const EigenSparseMatrix<Real_t> &, const EigenVector<Real_t> &);

    template ALGEBRA_EXPORT Real_t quadratic<Real_t> (const EigenSparseMatrix<Real_t> &, const EigenVector<Real_t> &);

    template ALGEBRA_EXPORT std::ostream &operator<< <Real_t> (std::ostream &s, const EigenVector<Real_t> &v);
    template ALGEBRA_EXPORT std::ostream &operator<< <Real_t> (std::ostream &s, const EigenDenseMatrix<Real_t> &m);
    template ALGEBRA_EXPORT std::ostream &operator<< <Real_t> (std::ostream &s, const EigenSparseMatrix<Real_t> &m);

    //------------------------------------------------------------------
    // Complex_t
    template class ALGEBRA_EXPORT EigenVector<Complex_t>;
    template class ALGEBRA_EXPORT EigenDenseMatrix<Complex_t>;
    template class ALGEBRA_EXPORT EigenSparseMatrix<Complex_t>;

    template ALGEBRA_EXPORT           EigenVector<Complex_t> operator+ <Complex_t> (const EigenVector<Complex_t> &x,
                                                                             const EigenVector<Complex_t> &y);
    template ALGEBRA_EXPORT           EigenVector<Complex_t> operator- <Complex_t> (const EigenVector<Complex_t> &,
                                                                             const EigenVector<Complex_t> &);
    template ALGEBRA_EXPORT           EigenVector<Complex_t> operator* <Complex_t> (const Complex_t &s, const EigenVector<Complex_t> &x);
    template ALGEBRA_EXPORT Complex_t operator* <Complex_t> (const EigenVector<Complex_t> &x, const EigenVector<Complex_t> &y);
    template ALGEBRA_EXPORT           EigenVector<Complex_t> operator/ <Complex_t> (const EigenVector<Complex_t> &x, const Complex_t &s);

    template ALGEBRA_EXPORT EigenSparseMatrix<Complex_t> operator+ <Complex_t> (const EigenSparseMatrix<Complex_t> &,
                                                                                   const EigenSparseMatrix<Complex_t> &);
    template ALGEBRA_EXPORT EigenSparseMatrix<Complex_t> operator- <Complex_t> (const EigenSparseMatrix<Complex_t> &,
                                                                                   const EigenSparseMatrix<Complex_t> &);
    template ALGEBRA_EXPORT EigenSparseMatrix<Complex_t> operator* <Complex_t> (const Complex_t &, const EigenSparseMatrix<Complex_t> &);
    template ALGEBRA_EXPORT EigenVector<Complex_t> operator* <Complex_t> (const EigenSparseMatrix<Complex_t> &,
                                                                             const EigenVector<Complex_t> &);

    template ALGEBRA_EXPORT Complex_t quadratic<Complex_t> (const EigenSparseMatrix<Complex_t> &, const EigenVector<Complex_t> &);

    template ALGEBRA_EXPORT std::ostream &operator<< <Complex_t> (std::ostream &s, const EigenVector<Complex_t> &v);
    template ALGEBRA_EXPORT std::ostream &operator<< <Complex_t> (std::ostream &s, const EigenDenseMatrix<Complex_t> &m);
    template ALGEBRA_EXPORT std::ostream &operator<< <Complex_t> (std::ostream &s, const EigenSparseMatrix<Complex_t> &m);

    template <typename Scalar, typename TripletIteratorType>
    void convertTripletsToMKLSparse(const TripletIteratorType& begin, const TripletIteratorType& end, Eigen::SparseMatrix<Scalar>& mat){

        // 1. Prepare COO format data 
        std::vector<MKL_INT> rows, cols;
        std::vector<Scalar> values;
        size_t nnz = std::distance(begin, end);
        if (0 == nnz) return;
        rows.reserve(nnz);
        cols.reserve(nnz);
        values.reserve(nnz);

        MKL_INT m = mat.rows();
        MKL_INT n = mat.cols();

        for (auto it = begin; it != end; ++it) {
            rows.push_back(static_cast<MKL_INT>(it->row()));
            cols.push_back(static_cast<MKL_INT>(it->col()));
            values.push_back(it->value());
        }

        // 2. Call the MKL interface to create a COO matrix
        sparse_matrix_t mklMat;
        sparse_status_t status;
        if constexpr (std::is_same_v<Scalar, float>) {
            status = mkl_sparse_s_create_coo(&mklMat, SPARSE_INDEX_BASE_ZERO,
                                        m, n, nnz,
                                        rows.data(), cols.data(), values.data());
        } else if constexpr (std::is_same_v<Scalar, double>) {
            status = mkl_sparse_d_create_coo(&mklMat, SPARSE_INDEX_BASE_ZERO,
                                        m, n, nnz,
                                        rows.data(), cols.data(), values.data());
        } else {
            static_assert(std::is_same_v<Scalar, float> || std::is_same_v<Scalar, double>,
                        "Only float/double are supported");
        }
        if (status != SPARSE_STATUS_SUCCESS) return;
        

        // 3. Convert to CSC format (used internally by Eigen)
        sparse_matrix_t csrMat;
        status = mkl_sparse_convert_csr(mklMat, SPARSE_OPERATION_TRANSPOSE , &csrMat);
        mkl_sparse_order(csrMat);
        assert(status == SPARSE_STATUS_SUCCESS);
        
        // 4. Get CSR data pointer
        sparse_index_base_t indexing;
        MKL_INT csrRows, csrCols;
        MKL_INT *csrStart, *csrEnd, *csrInd;
        Scalar *csrValues;
        if constexpr (std::is_same_v<Scalar, float>) {
            status = mkl_sparse_s_export_csr(csrMat, &indexing, &csrRows, &csrCols,
                                        &csrStart, &csrEnd, &csrInd, &csrValues);
        } else if constexpr (std::is_same_v<Scalar, double>){
            status = mkl_sparse_d_export_csr(csrMat, &indexing, &csrRows, &csrCols,
                                        &csrStart, &csrEnd, &csrInd, &csrValues);
        }else {
            static_assert(std::is_same_v<Scalar, float> || std::is_same_v<Scalar, double>,
                        "Only float/double are supported");
        }
        assert(status == SPARSE_STATUS_SUCCESS);

        mat.resize(csrCols, csrRows);
        size_t actual_nnz = csrStart[csrRows] - csrStart[0];
        mat.reserve(actual_nnz);
        mat.makeCompressed();

        // Eigen::VectorXi colSize(csrRows);
        // for(auto i = 1; i <= csrRows; ++i){
        //     colSize[i - 1] = csrStart[i] - csrStart[i - 1];
        // }
        // mat.reserve(colSize);
        mat.makeCompressed();
        if (mat.outerIndexPtr() && mat.innerIndexPtr() && mat.valuePtr()) {
            mat.data().resize(actual_nnz);
            std::memcpy(mat.outerIndexPtr(), csrStart, (csrRows + 1) * sizeof(MKL_INT));
            std::memcpy(mat.innerIndexPtr(), csrInd, actual_nnz * sizeof(MKL_INT));
            std::memcpy(mat.valuePtr(), csrValues, actual_nnz * sizeof(Scalar));
        }

        // 6. 释放MKL资源
        mkl_sparse_destroy(mklMat);
        mkl_sparse_destroy(csrMat);

    }

    template ALGEBRA_EXPORT void convertTripletsToMKLSparse<double> (const std::vector<Eigen::Triplet<double>>::iterator &,
                                                                     const std::vector<Eigen::Triplet<double>>::iterator &,
                                                                     Eigen::SparseMatrix<double> &);

    template ALGEBRA_EXPORT void convertTripletsToMKLSparse<float> (const std::vector<Eigen::Triplet<float>>::iterator &,
                                                                    const std::vector<Eigen::Triplet<float>>::iterator &,
                                                                    Eigen::SparseMatrix<float> &);
}  // namespace Algebra
SG_NAMESPACE_CLOSE